package org.example.function;

import com.antgroup.geaflow.api.context.RuntimeContext;
import com.antgroup.geaflow.api.function.RichFunction;
import com.antgroup.geaflow.api.function.io.SourceFunction;
import com.antgroup.geaflow.api.window.IWindow;
import org.example.csv.CSVData;
import org.example.util.CSVFaster;

import java.io.Serializable;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Objects;

public class CSVFileSource<OUT> extends RichFunction implements SourceFunction<OUT> {

    private final Path filePath;

    private final List<String> selectedHeaders;

    private CSVData csvData;

    private final CSVLineParser<OUT> parser;

    private Integer readPos = null;

    private List<OUT> records;

    private List<String[]> lines;

    public CSVFileSource(Path filePath, List<String> selectedHeaders, CSVLineParser<OUT> parser) {
        this.filePath = filePath;
        this.selectedHeaders = selectedHeaders;
        this.parser = parser;
    }

    @Override
    public void open(RuntimeContext runtimeContext) {
    }

    @Override
    public void close() {

    }

    @Override
    public void init(int parallel, int index) {
        // TODO source parallelism: ref FileSource ?
        // long st = System.currentTimeMillis();
        this.csvData = CSVFaster.read(this.filePath, selectedHeaders);
        // this.csvData = CSVFaster2.read(Paths.get(this.filePath), selectedHeaders);
        // System.out.println("read cost: " + (System.currentTimeMillis() - st) + " ms");
        // this.records = this.parseLines();
        this.lines = this.csvData.getLines();
    }

    private List<OUT> parseLines() {
        List<OUT> result = new ArrayList<>((int) this.csvData.count());
        List<String[]> lines = this.csvData.getLines();
        for (String[] line : lines) {
            Collection<OUT> collection = this.parser.parse(Arrays.asList(line));
            result.addAll(collection);
        }
        return result;
    }

    @Override
    public boolean fetch(IWindow<OUT> window, SourceContext<OUT> ctx) throws Exception {
        if (Objects.isNull(readPos)) {
            readPos = 0;
        }

        while (readPos < this.lines.size()) {
            Collection<OUT> outs = this.parser.parse(Arrays.asList(this.lines.get(readPos)));
            for (OUT out : outs) {
                long windowId = window.assignWindow(out);
                if (windowId == window.windowId()) {
                    ctx.collect(out);
                } else {
                    break;
                }
            }
            readPos++;
        }

        return readPos <= this.lines.size();
    }

    public interface CSVLineParser<OUT> extends Serializable {

        Collection<OUT> parse(List<String> line);
    }
}
